package com.lili.collection.core.context.impl;

import com.lili.collection.core.context.ExchangeStageContext;
import com.lili.collection.core.context.StageContext;
import com.lili.collection.core.runner.StageRunner;

import java.util.HashMap;
import java.util.Map;

/**
 * 单值交换上下文
 *
 * @author lili
 * @date 2023/3/9 0:41
 */
public class SingleValueExchangeStageContext implements ExchangeStageContext {

    private Map<StageRunner.StageType, StageContext> stageContextMap;

    public SingleValueExchangeStageContext() {
        this.stageContextMap = new HashMap<>();
        for (StageRunner.StageType stageType : StageRunner.StageType.values()) {
            this.stageContextMap.put(stageType, null);
        }
    }

    @Override
    public StageContext take(StageRunner.StageType stageType) {
        StageContext stageContext = stageContextMap.get(stageType);
        stageContextMap.put(stageType, null);
        return stageContext;
    }

    @Override
    public void put(StageRunner.StageType stageType, StageContext stageContext) {
        this.stageContextMap.put(stageType, stageContext);
    }

    @Override
    public StageContext peek(StageRunner.StageType stageType) {
        return stageContextMap.get(stageType);
    }

    @Override
    public void clear() {
        for (StageRunner.StageType stageType : StageRunner.StageType.values()) {
            this.stageContextMap.put(stageType, null);
        }
    }
}
